#!/usr/bin/env python3

import argparse
import os
from pathlib import Path
import subprocess
import re


HOME = os.getenv("HOME")
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.realpath(os.path.abspath(__file__))))
CUID, CGID = os.getuid(), os.getgid()
LOGDEVICE_CONTAINER_NAME = f"logdevice-dev-cluster-{CUID}"
DEFAULT_LOGDEVICE_DIR = os.path.join(PROJECT_ROOT, "local-data/logdevice")

Path(DEFAULT_LOGDEVICE_DIR).mkdir(parents=True, exist_ok=True)

logerr = lambda s: print(f"\033[91m{s}\033[0m")
logdebug = lambda s: print(f"\033[95m[DEBUG] \033[0m{s}")
loginfo = lambda s: print(f"\033[96m{s}\033[0m")


def run_sh(sh, stderr=subprocess.STDOUT, stdout=None, check=True):
    assert isinstance(sh, str)
    args = ["bash", "-c", sh]
    return subprocess.run(args, stderr=stderr, stdout=stdout, check=check)


def is_cluster_started(container_bin):
    result = run_sh(
        str(container_bin) + " ps --format {{.Names}}",
        stdout=subprocess.PIPE
    )
    if result and result.stdout:
        rs = result.stdout.decode().strip().split('\n')
        if LOGDEVICE_CONTAINER_NAME in rs:
            return True
    return False


def cluster_start(image, data_dir, container_bin):
    if is_cluster_started(container_bin):
        return logerr("Already started!")

    cmd = ' '.join([
        f'{container_bin} run -td --name {LOGDEVICE_CONTAINER_NAME} --rm',
        '--network host',
        '-u' + str(CUID) + ":" + str(CGID) if container_bin.strip() in ["docker"] else '',
        f'-v {data_dir}:/data/store {image}',
        '/usr/local/bin/ld-dev-cluster --root /data/store --use-tcp'
    ])
    run_sh(cmd)


def cluster_stop(container_bin):
    run_sh(f"{container_bin} kill {LOGDEVICE_CONTAINER_NAME}", check=False)


def run_ldshell(image, data_dir, container_bin):
    if not is_cluster_started(container_bin):
        return logerr("Please start the dev-cluster first. Run: dev-tools cluster-start")

    result = run_sh(f"{container_bin} logs {LOGDEVICE_CONTAINER_NAME}", stdout=subprocess.PIPE)
    if result and result.stdout:
        rs = set(re.findall(r'--admin-server-port=(\d*)', result.stdout.decode()))
        if len(rs) != 1:
            return logerr("Get admin-server-port error!")
        host = '127.0.0.1'
        port = rs.pop()
        cmd = ' '.join([
            f'{container_bin} run -it --rm --network host {image}',
            f'/usr/local/bin/ldshell --admin-server-host={host} --admin-server-port={port}'
        ])
        run_sh(cmd, check=False)


def haskell_env(container_bin, image, command, cmd_args, data_dir,
                tty, interactive, rm, extra_container_options):

    if not is_cluster_started(container_bin):
        return logerr("Please start the dev-cluster first. Run: dev-tools cluster-start")

    stack_home = f"{HOME}/.stack"
    cabal_home = f"{HOME}/.cabal"
    hoogle_home = f'{HOME}/.hoogle'

    Path(f"{HOME}/.local/bin").mkdir(parents=True, exist_ok=True)
    Path(f"{HOME}/.ghc").mkdir(parents=True, exist_ok=True)
    Path(f"{cabal_home}").mkdir(parents=True, exist_ok=True)
    Path(f"{stack_home}").mkdir(parents=True, exist_ok=True)
    Path(f"{hoogle_home}").mkdir(parents=True, exist_ok=True)

    user_cmd = f"{command} {' '.join(cmd_args)}"
    user_path = f"/opt/ghc/bin:/opt/cabal/bin:{HOME}/.local/bin:{HOME}/.cabal/bin:$PATH"
    container_opts = " ".join([
        "-t" if tty else "",
        "-i" if interactive else "",
        "--rm" if rm else "",
        "-u " + str(CUID) + ":" + str(CGID) if container_bin.strip() in ["docker"] else ""
    ])
    cmd = f'''{container_bin} run {container_opts} \
                -e HOME={HOME} \
                -e LC_ALL=C.UTF-8 \
                -e PATH={user_path} \
                -v {HOME}/.local/bin:{HOME}/.local/bin:rw \
                -v {HOME}/.ghc:{HOME}/.ghc:rw \
                -v {stack_home}:{HOME}/.stack:rw \
                -v {cabal_home}:{HOME}/.cabal:rw \
                -v {hoogle_home}:{HOME}/.hoogle:rw \
                -v "{PROJECT_ROOT}:{PROJECT_ROOT}" \
                -v "{data_dir}:/data/store" \
                -v "/tmp:/tmp" \
                -w "{PROJECT_ROOT}" \
                --network host \
                {extra_container_options} \
                {image} {user_cmd}'''
    loginfo(f"Run <{user_cmd}> from <{image}> image...")
    run_sh(cmd, check=False)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='HStream dev tools.')
    subparsers = parser.add_subparsers(title='subcommands', dest='sub_command')

    parser_cluster_start = subparsers.add_parser(
        'cluster-start', help='Start logdevice dev cluster')
    parser_cluster_start.add_argument(
        '--image', '-i', help='logdevice docker images',
        default='docker.io/hstreamdb/logdevice')
    parser_cluster_start.add_argument(
        '--data-dir', '-d', type=Path, default=DEFAULT_LOGDEVICE_DIR)
    parser_cluster_start.add_argument('--container-bin', default='docker')

    parser_cluster_stop = subparsers.add_parser(
        'cluster-stop', help='Stop logdevice dev cluster')
    parser_cluster_stop.add_argument('--container-bin', default='docker')

    parser_cluster_ldshell = subparsers.add_parser(
        'ldshell', help='')
    parser_cluster_ldshell.add_argument(
        '--image', '-i', help='logdevice docker images',
        default='docker.io/hstreamdb/logdevice')
    parser_cluster_ldshell.add_argument('--container-bin', default='docker')
    parser_cluster_ldshell.add_argument(
        '--data-dir', '-d', type=Path, default=DEFAULT_LOGDEVICE_DIR)

    parser_haskell = subparsers.add_parser(
        'shell', help='Enter in a shell with all haskell dev dependencies')
    parser_haskell.add_argument(
        '--image', '-i', help='logdevice docker images', default='docker.io/hstreamdb/haskell')
    parser_haskell.add_argument('--container-bin', default='docker')
    parser_haskell.add_argument(
        '--data-dir', '-d', type=Path, default=DEFAULT_LOGDEVICE_DIR)
    parser_haskell.add_argument('--extra-container-options', default='')

    parser_haskell_cabal = subparsers.add_parser('cabal', help='Run cabal command directly')
    parser_haskell_cabal.add_argument(
        '--image', '-i', help='logdevice docker images', default='docker.io/hstreamdb/haskell')
    # TODO: since python3.9 there is BooleanOptionalAction available in argparse
    parser_haskell_cabal.add_argument('--no-interactive', action='store_true')
    parser_haskell_cabal.add_argument('--no-tty', action='store_true')
    parser_haskell_cabal.add_argument('--container-bin', default='docker')
    parser_haskell_cabal.add_argument(
        '--data-dir', '-d', type=Path, default=DEFAULT_LOGDEVICE_DIR)
    parser_haskell_cabal.add_argument('--extra-container-options', default='')
    parser_haskell_cabal.add_argument('cmd_args', nargs='*')

    args = vars(parser.parse_args())
    sub_command = args.pop('sub_command')

    if sub_command == 'cluster-start':
        cluster_start(**args)
    elif sub_command == 'ldshell':
        run_ldshell(**args)
    elif sub_command == 'cluster-stop':
        cluster_stop(**args)
    elif sub_command == 'shell':
        haskell_env(command='bash', cmd_args="", tty=True, interactive=True, rm=True, **args)
    elif sub_command == 'cabal':
        args['interactive'] = not args.pop('no_interactive')
        args['tty'] = not args.pop('no_tty')
        haskell_env(command='/opt/ghc/bin/cabal', rm=True, **args)
    else:
        parser.print_help()
